#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2018 crane <crane@his-pc>
#
# Distributed under terms of the MIT license.

"""
https://www.lintcode.com/problem/sliding-window-unique-elements-sum/
"""

from collections import Counter


class Solution:
    """
    @param nums: the given array
    @param k: the window size
    @return: the sum of the count of unique elements in each window
    """
    def slidingWindowUniqueElementsSum(self, nums, k):
        self.nums_len = len(nums)
        self.nums = nums
        self.k = k

        self.counter = Counter()
        self.sum = 0
        self.count = 0

        if self.nums_len < k:
            k = self.k = self.nums_len


        # 计算第一个窗口
        self.count_first_window()
        self.sum += self.count

        for self.right in range(k, self.nums_len):
            # 计算滑动窗口
            self.count_slide_window()
            self.sum += self.count

        return self.sum

    def count_first_window(self):
        for i in range(self.k):
            self.slide_in(i)

    def count_slide_window(self):
        self.old_left = self.right - self.k

        self.slide_out(self.old_left)
        self.slide_in(self.right)

    def slide_out(self, idx):
        # truncate left
        n = self.nums[idx]
        self.counter[n] -= 1
        v = self.counter[n]
        assert v >= 0

        if v == 1:
            self.count += 1
        elif v == 0:
            self.count -= 1

    def slide_in(self, idx):
        # extend right
        n = self.nums[idx]
        self.counter[n] += 1
        v = self.counter[n]
        assert v >= 1

        if v == 1:
            self.count += 1
        elif v == 2:
            self.count -= 1


def test():
    s = Solution()
    # sum = s.slidingWindowUniqueElementsSum([1, 2, 1, 3, 3], 3)
    # sum = s.slidingWindowUniqueElementsSum([1,1,1,1,1] , 2)
    l = [27,14,60,87,37,53,100,18,51,37,14,57,22,95,50,83,41,43,36,48,52,97,16,46,75,24,47,13,40,40,48,45,56,58,77,3,78,60,31,27,40,53,57,29,30,65,37,77,1,40,89,100,50,49,100,51,22,66,33,33,70,36,64,70,11,27,57,77,17,28,62,70,32,88,12,47,69,30,93,3,47,69,64,88,7,40,38,5,23,4,58,97,19,55,17,23]

    sum = s.slidingWindowUniqueElementsSum(l, 21)
    print('sum is [%s]' % sum)

def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
